# Copyright (c) Meta Platforms, Inc. and affiliates.

class StreamType(Enum):
    serialized: int
    numeric: int
    fixed_size_field: int
    variable_size_field: int

class Graph:
    def input_type(self) -> StreamType: ...

class JsonGraph(Graph):
    def __init__(
        self,
        json: str | Dict,
        input_type: StreamType = StreamType.serialized,
        custom_transforms: Dict | None = None,
        custom_graphs: Dict | None = None,
        custom_selectors: Dict | None = None,
    ) -> None: ...

def compress(
    data: bytes,
    graph: Graph,
    global_params: Dict[GCParam, int] | None = None,
) -> bytes: ...
def compress_multi(
    data: List[bytes],
    graph: Graph,
    global_params: Dict[GCParam, int] | None = None,
) -> bytes: ...
def decompress(
    compressed: bytes,
    graph: Graph | None = None,
) -> bytes: ...
def decompress_multi(
    compressed: bytes,
    graph: Graph | None = None,
) -> List[bytes]: ...
def get_header_size(compressed: bytes) -> int: ...
def measure_decompress_speed_one(
    compressed: bytes,
    graph: Graph | None = None,
) -> float: ...
def measure_decompress_speed_multiple(
    compressed: bytes,
    graph: Graph | None = None,
) -> float: ...
def split_extracted_streams(
    streams: bytes,
) -> np.array: ...
def read_extracted_streams(
    path: str,
) -> np.array: ...

class GCParam(Enum):
    compression_level: int
    decompression_level: int
    format_version: int

class Stream:
    def type(self) -> StreamType: ...
    def as_array(self) -> np.array: ...
    def as_bytes(self) -> bytes: ...
    def as_list(self) -> List[bytes]: ...

class EncoderCtx:
    def get_global_param(self, param: GCParam) -> int: ...
    def get_local_int_param(self, param: GCParam) -> int: ...
    def get_local_string_param(self, param: GCParam) -> str: ...
    def get_local_binary_param(self, param: GCParam) -> bytes: ...
    def send_transform_header(self, header: bytes) -> None: ...
    def get_inputs(self) -> List[Stream]: ...
    def get_input(self, idx: int) -> Stream: ...
    def create_output(self, idx: int, type: StreamType) -> Stream: ...

class SimpleEncoderCtx(EncoderCtx):
    def get_global_param(self, param: GCParam) -> int: ...
    def get_local_int_param(self, param: GCParam) -> int: ...
    def get_local_string_param(self, param: GCParam) -> str: ...
    def get_local_binary_param(self, param: GCParam) -> bytes: ...
    def send_transform_header(self, header: bytes) -> None: ...

class SelectorCtx:
    def get_global_param(self, param: GCParam) -> int: ...
    def get_local_int_param(self, param: GCParam) -> int: ...
    def get_local_string_param(self, param: GCParam) -> str: ...
    def get_local_binary_param(self, param: GCParam) -> bytes: ...
    def is_graph_supported(self, graph: Graph) -> bool: ...
    def try_graph(self, graph: Graph) -> bool: ...

class DecoderCtx:
    def get_transform_header(self) -> bytes: ...
    def get_fixed_inputs(self) -> List[Stream]: ...
    def get_variable_inputs(self) -> List[Stream]: ...
    def create_output(self, type: StreamType) -> Stream: ...

class SimpleDecoderCtx(DecoderCtx):
    def get_transform_header(self) -> bytes: ...

class CustomTransform:
    def __init__(
        self,
        id: int,
        input_type: List[StreamType],
        fixed_output_types: List[StreamType],
        variable_output_types: List[StreamType],
        docs: str = "",
    ) -> None: ...
    def encode(self, ctx: EncoderCtx, input: Stream) -> List[Stream]: ...
    def decode(self, ctx: DecoderCtx, inputs: List[Stream]) -> Stream: ...

class SimpleCustomTransform(CustomTransform):
    def __init__(
        self,
        id: int,
        input_type: StreamType,
        output_types: List[StreamType],
        docs: str = "",
    ) -> None: ...

class CustomSelector:
    def __init__(
        self,
        input_type: StreamType,
        docs: str = "",
    ) -> None: ...
    def select(
        self,
        ctx: SelectorCtx,
        input: Stream,
        graphs: List[Graph],
    ) -> Graph: ...
